from numpy import (
    ndarray,
    pi,
)

from src.utils.real_complex_vector_reshaping import (
    real_vector_to_half_complex_vector,
    complex_vector_to_double_real_vector,
    rad_and_phase_to_complex_vector,
    complex_vector_to_rad_and_phase,
)


# TODO: 硬编码的规范值可能会导致在不同应用中出现问题

def get_state_erroneous_channel_state_information(
        satellites,
        csi_format: str,
        norm_csi: bool,
) -> ndarray:
    """
    获取卫星的错误的信道状态信息，并根据指定格式进行处理和归一化。

    参数:
        satellites: 包含卫星数据的对象，要求有属性 `erroneous_channel_state_information`
        csi_format: CSI（信道状态信息）的格式，可以是 'rad_phase'（弧度和相位）或 'real_imag'（实部和虚部）
        norm_csi: 是否对CSI进行归一化处理

    返回:
        ndarray: 处理后的CSI状态信息
    """
    erroneous_csi = satellites.erroneous_channel_state_information.flatten()  # 将错误的信道状态信息展平为一维数组

    if csi_format == 'rad_phase':
        # 如果CSI格式为弧度和相位，将其转换为弧度和相位表示
        state_real = complex_vector_to_rad_and_phase(erroneous_csi)
        if norm_csi:
            # 归一化处理
            half_length_idx = int(len(state_real) / 2)  # 获取数组的一半长度
            state_real[:half_length_idx] = state_real[:half_length_idx] * 1e7  # 将弧度部分乘以1e7
            state_real[half_length_idx:] = state_real[half_length_idx:] / pi  # 将相位部分除以π

    elif csi_format == 'real_imag':
        # 如果CSI格式为实部和虚部，将其转换为双重实部表示
        state_real = complex_vector_to_double_real_vector(erroneous_csi)
        if norm_csi:
            # 归一化处理
            state_real *= 1e8  # 将实部和虚部都乘以1e8

    else:
        # 如果CSI格式不符合要求，则抛出异常
        raise ValueError(f'Unknown CSI Format {csi_format}')

    return state_real  # 返回处理后的CSI状态信息


def get_state_aods(
        satellites,
) -> ndarray:
    """
    获取卫星到用户的到达方向（AOD）。

    参数:
        satellites: 包含卫星数据的对象，要求有方法 `get_aods_to_users`

    返回:
        ndarray: 卫星到用户的到达方向，展平为一维数组
    """
    state = satellites.get_aods_to_users()  # 获取卫星到用户的到达方向
    return state.flatten()  # 将结果展平为一维数组并返回
